import os
from torch.utils import data
from torchvision import transforms
from PIL import Image


class ImageFolder(data.Dataset):
    def __init__(self, dir, transform=None):
        self.img_paths = list(map(lambda x: os.path.join(dir, x), os.listdir(dir)))
        self.transform = transform

    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        img = Image.open(img_path).convert('RGB')
        if self.transform:
            img = self.transform(img)

        return img

    def __len__(self):
        return len(self.img_paths)


def get_loader(img_path, img_size, batch_size, num_workers=2):
    transform = transforms.Compose(
        [
            transforms.Scale(img_size),
            transforms.ToTensor(),
            transforms.Normalize((.5, .5, .5), (.5, .5, .5))
        ]
    )

    dataset = ImageFolder(img_path, transform)
    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batch_size,
                                  shuffle=True,
                                  num_workers=num_workers)
    return data_loader